// BaseSegDetector.hpp
#ifndef YOLOS_EDGEPLATFORM_BASE_SEG_DETECTOR_HPP
#define YOLOS_EDGEPLATFORM_BASE_SEG_DETECTOR_HPP

/**
 * @file BaseSegDetector.hpp
 * @brief 边缘平台 YOLO 实例分割检测器的抽象基类
 *
 * @author FANKYT
 * @date 2025
 */

#include <string>
#include <vector>
#include <memory>
#include <opencv2/opencv.hpp>

#include "tools/Common.hpp"
#include "tools/ScopedTimer.hpp"

namespace yolos_edgeplatform {

/**
 * @brief YOLO 实例分割检测器的抽象基类
 */
class BaseSegDetector {
public:
    virtual ~BaseSegDetector() = default;

    /**
     * @brief 对提供的图像执行实例分割检测
     *
     * @param image 输入 OpenCV Mat 图像 (BGR 格式)
     * @param confThreshold 置信度阈值 (默认 0.25)
     * @param nmsThreshold NMS IoU 阈值 (默认 0.7)
     * @param maskThreshold 掩码二值化阈值 (默认 0.5)
     * @return std::vector<Segmentation> 分割检测结果向量
     */
    virtual std::vector<Segmentation> detect(const cv::Mat &image,
                                              float confThreshold = 0.25f,
                                              float nmsThreshold = 0.7f,
                                              float maskThreshold = 0.5f) = 0;

    /**
     * @brief 获取模型期望的输入尺寸
     */
    virtual cv::Size getInputSize() const = 0;

    /**
     * @brief 获取模型支持的类别数量
     */
    virtual int getNumClasses() const = 0;

    /**
     * @brief 获取类别名称
     */
    virtual const std::vector<std::string>& getClassNames() const = 0;

    /**
     * @brief 设置预处理方法
     */
    virtual void setPreprocessType(PreprocessType type) = 0;

protected:
    /**
     * @brief 从文件加载类别名称
     */
    std::vector<std::string> loadClassNames(const std::string &path) {
        std::vector<std::string> classNames;
        std::ifstream file(path);
        if (!file.is_open()) {
            throw std::runtime_error("[ERROR] Cannot open class names file: " + path);
        }

        std::string line;
        while (std::getline(file, line)) {
            if (!line.empty()) {
                classNames.push_back(line);
            }
        }

        return classNames;
    }

    /**
     * @brief LetterBox 预处理 - 保持宽高比缩放并填充
     */
    cv::Mat letterBox(const cv::Mat &image, cv::Size targetSize,
                      float &scaleX, float &scaleY, int &padX, int &padY) {
        int origW = image.cols;
        int origH = image.rows;
        int targetW = targetSize.width;
        int targetH = targetSize.height;

        float scale = std::min(float(targetW) / origW, float(targetH) / origH);
        scaleX = scaleY = scale;

        int newW = static_cast<int>(origW * scale);
        int newH = static_cast<int>(origH * scale);

        cv::Mat resized;
        cv::resize(image, resized, cv::Size(newW, newH));

        padX = (targetW - newW) / 2;
        padY = (targetH - newH) / 2;

        cv::Mat output;
        cv::copyMakeBorder(resized, output, padY, targetH - newH - padY,
                          padX, targetW - newW - padX,
                          cv::BORDER_CONSTANT, cv::Scalar(114, 114, 114));

        return output;
    }

    /**
     * @brief Resize 预处理 - 直接拉伸到目标尺寸
     */
    cv::Mat resizeImage(const cv::Mat &image, cv::Size targetSize,
                        float &scaleX, float &scaleY) {
        scaleX = float(targetSize.width) / image.cols;
        scaleY = float(targetSize.height) / image.rows;

        cv::Mat resized;
        cv::resize(image, resized, targetSize);
        return resized;
    }
};

} // namespace yolos_edgeplatform

#endif // YOLOS_EDGEPLATFORM_BASE_SEG_DETECTOR_HPP
